# EBML/Matroska parser
# Copyright (C) 2010  Harry Bragg <h.bragg@gmail.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
#

# this code is based on the ebml parser by Johannes Sasongko 
# <sasongko@gmail.com> obtained from exaile repository at revision 3765 
# (2010-11-12), under /xl/metadata/_matroska.py which in turn is heavily 
# based on public domain code by "Omion" (from the Hydrogenaudio 
# forums), as obtained from Matroska's Subversion repository at
# revision 858 (2004-10-03), under "/trunk/Perl.Parser/MatroskaParser.pm".


import sys
from struct import unpack, pack
import os

SINT, UINT, FLOAT, STRING, UTF8, DATE, MASTER, BINARY = range(8)

EBML_ID_VOID = 0x6A
EBML_ID_CRC32 = 0xBF

EBML_SIZE_UNKNOWN = 0x00ffffffffffffff

class EbmlException(Exception): pass

class BinaryData(str): pass
class UnknownData: pass

class EbmlReader:
    """EBML parser.

    Usage: Ebml(location, tags).parse()
    tags is a dictionary of the form { id: (name, type) }.
    """

    ## Constructor and destructor

    def __init__(self, location, tags):
        self.tags = tags
        self.open(location)
        self.storePosition = False

    def __del__(self):
        self.close()

    ## File access.
    ## These can be overridden to provide network support.

    def open(self, location):
        """Open a location and set self.size."""
        self.file = f = open(location, 'rb')
        f.seek(0, os.SEEK_END)
        self.size = f.tell()
        f.seek(0, os.SEEK_SET)

    def seek(self, offset, mode = os.SEEK_SET):
        self.file.seek(offset, mode)

    def tell(self):
        return self.file.tell()

    def read(self, length):
        return self.file.read(length)

    def close(self):
        self.file.close()

    ## Element reading

    def readSize(self):
        b1 = self.read(1)
        b1b = ord(b1)
        if b1b & 0x80:
            # 1 byte
            return b1b & 0x7f
        elif b1b & 0x40:
            # 2 bytes
            # JS: BE-ushort
            return unpack(">H", chr(0x40 ^ b1b) + self.read(1))[0]
        elif b1b & 0x20:
            # 3 bytes
            # JS: BE-ulong
            return unpack(">L", "\0" + chr(0x20 ^ b1b) + self.read(2))[0]
        elif b1b & 0x10:
            # 4 bytes
            # JS: BE-ulong
            return unpack(">L", chr(0x10 ^ b1b) + self.read(3))[0]
        elif b1b & 0x08:
            # 5 bytes
            # JS: uchar BE-ulong. We change this to BE uchar ulong.
            high, low = unpack(">BL", chr(0x08 ^ b1b) + self.read(4))
            return high * 4294967296 + low
        elif b1b & 0x04:
            # 6 bytes
            # JS: BE-slong BE-ulong
            high, low = unpack(">HL", chr(0x04 ^ b1b) + self.read(5))
            return high * 4294967296 + low
        elif b1b & 0x02:
            # 7 bytes
            # JS: BE-ulong BE-ulong
            high, low = unpack(">LL",
                    "\0" + chr(0x02 ^ b1b) + self.read(6))
            return high * 4294967296 + low
        elif b1b & 0x01:
            # 8 bytes
            # JS: BE-ulong BE-ulong
            high, low = unpack(">LL", chr(0x01 ^ b1b) + self.read(7))
            return high * 4294967296 + low
        else:
            raise EbmlException(
                    "invalid element size with leading byte 0x%X" % b1b)

    def readInteger(self, length):
        if length == 1:
            # 1 byte
            return ord(self.read(1))
        elif length == 2:
            # 2 bytes
            return unpack(">H", self.read(2))[0]
        elif length == 3:
            # 3 bytes
            return unpack(">L", "\0" + self.read(3))[0]
        elif length == 4:
            # 4 bytes
            return unpack(">L", self.read(4))[0]
        elif length == 5:
            # 5 bytes
            high, low = unpack(">BL", self.read(5))
            return high * 4294967296 + low
        elif length == 6:
            # 6 bytes
            high, low = unpack(">HL", self.read(6))
            return high * 4294967296 + low
        elif length == 7:
            # 7 bytes
            high, low = unpack(">LL", "\0" + (self.read(7)))
            return high * 4294967296 + low
        elif length == 8:
            # 8 bytes
            high, low = unpack(">LL", self.read(8))
            return high * 4294967296 + low
        else:
            raise EbmlException(
                    "don't know how to read %r-byte integer" % length)

    def readFloat(self, length):
        # Need to reverse the bytes for little-endian machines
        if length == 4:
            # single
            return unpack('@f', self.read(4)[::-1])[0]
        elif length == 8:
            # double
            return unpack('@d', self.read(8)[::-1])[0]
        elif length == 10:
            # extended (don't know how to handle it)
            return 'EXTENDED'
        else:
            raise EbmlException("don't know how to read %r-byte float" % length)

    def readID(self):
        b1 = self.read(1)
        b1b = ord(b1)
        if b1b & 0x80:
            # 1 byte
            return b1b & 0x7f
        elif b1b & 0x40:
            # 2 bytes
            return unpack(">H", chr(0x40 ^ b1b) + self.read(1))[0]
        elif b1b & 0x20:
            # 3 bytes
            return unpack(">L", "\0" + chr(0x20 ^ b1b) + self.read(2))[0]
        elif b1b & 0x10:
            # 4 bytes
            return unpack(">L", chr(0x10 ^ b1b) + self.read(3))[0]
        else:
            raise EbmlException(
                    "invalid element ID with leading byte 0x%X" % b1b)

    ## Parsing

    def parse(self, from_=0, to=None):
        """Parses EBML from `from_` to `to`.

        Note that not all streams support seeking backwards, so prepare to handle
        an exception if you try to parse from arbitrary position.
        """
        if to is None:
            to = self.size
        self.seek(from_, 0)
        node = {}
        # Iterate over current node's children.
        while self.tell() < to:
            pos = self.tell()
            try:
                id = self.readID()
            except EbmlException, e:
                # Invalid EBML header. We can't reliably get any more data from
                # this level, so just return anything we have.
                print >>sys.stderr, "ERROR:", e
                return node
            size = self.readSize()
            
            # quit on certain identifiers
            try:
                key, type_ = self.quitOn[id]
                break
            except KeyError:                            
                try:
                    key, type_ = self.tags[id]
                except KeyError:
                    self.seek(size, 1)
                else:
                    try:
                        if type_ is MASTER:
                            tell = self.tell()
                            value = self.parse(tell, tell + size)
                        elif type_ in (SINT, UINT, DATE):
                            value = self.readInteger(size)
                        elif type_ is FLOAT:
                            value = self.readFloat(size)
                        elif type_ is STRING:
                            value = unicode(self.read(size), 'ascii')
                        elif type_ is UTF8:
                            value = unicode(self.read(size), 'utf-8')
                        elif type_ is BINARY:
                            value = BinaryData(self.read(size))
                        else:
                            assert False
                    except (EbmlException, UnicodeDecodeError), e:
                        print >>sys.stderr, "WARNING:", e
                    try:
                        parentval = node[key]
                    except KeyError:
                        parentval = node[key] = []
                    if self.storePosition:
                        parentval.append((value, pos, size))
                    else:
                        parentval.append(value)
        return node

class EbmlContainer:
    """EbmlContainer, contain information about a MASTER ebml group
    """

    def __init__(self, id):
        self.id = id
        self.position = 0
        self.dataPosition = 0
        self.data = ''
        
    def size(self):
        return len(self.data)
    

# embl writer

class EbmlWriter(EbmlReader):
    """EmblWriter
    """
    
    def __init__(self, location):
        self.open(location)
        self.useBuffer = True
        
    def __del__(self):
        self.close()
        
    # file io, override for network etc
        
    def open(self, location):
        self.file = f = open(location, 'rb+') # read-write, keeping the existing data
        f.seek(0, os.SEEK_END)
        self.size = f.tell()
        f.seek(0, os.SEEK_SET)
        self.containers = []
        
    def write(self, data):
        if self.useBuffer and len(self.containers) > 0:
            # if we are using buffers, peek at the top most item on the containers list
            # and append the data to its buffer
            c = self.containers[len(self.containers)-1]
            c.data += data
        else:            
            # write out the data immediately
            self.file.write(data)
    
    def flush(self):
        self.file.flush()
        
    def truncate(self):
        self.file.truncate()
        
    # helper functions  
    
    def _toWriteID(self, value, length, mask):
        r = self._toWriteInt(value, length)
        return '{}{}'.format(chr(mask | ord(r[0])), r[1:])
    
    def _toWriteSize(self, value, length, mask):
        return self._toWriteID(value, length, mask)
    
    def _toWriteInt(self, value, length):
        if length == 1:
            # 1 byte
            return chr(value)
        elif length == 2:
            # 2 bytes
            return pack(">H", value)
        elif length == 3:
            # 3 bytes
            h = value / 65536
            l = value % 65536
            return pack(">BH", h, l)
        elif length == 4:
            # 4 bytes
            return pack(">L", value)
        elif length == 5:
            # 5 bytes
            h = value / 4294967296
            l = value % 4294967296
            return pack(">BL", h, l)
        elif length == 6:
            # 6 bytes
            h = value / 4294967296
            l = value % 4294967296            
            return pack(">HL", h, l)
        elif length == 7:
            # 7 bytes
            h = value / 4294967296
            h1 = h / 65536
            h2 = h % 65536
            l = value % 4294967296            
            return pack(">BHL", h1, h2, l)
        elif length == 8:
            # 8 bytes
            h = value / 4294967296
            l = value % 4294967296    
            return pack(">LL", h, l)
        else:
            raise EbmlException(
                    "don't know how to write %r-byte integer" % length)
            
    def _toWriteFloat(self, value, length):
        if length == 4:
            # single
            return pack('@f', value)
        elif length == 8:
            # double
            return pack('@d', value)
        else:
            raise EbmlException(
                    "don't know how to write %r-byte float" % length)
                    
    
    def _writeID(self, id):
        bytes = 1
        mask = 0x80
        
        # get the ID length
        while ((id >> ((bytes - 1) * 8)) >= (mask - 1) and bytes <= 4):
            mask >>= 1
            bytes += 1
            
        # if invalid ID, use dummy
        if bytes > 4:
            # warn: invalid ID
            bytes = 1
            id = EBML_ID_VOID
        
        # write
        try:
            self.write(self._toWriteID(id, bytes, mask))
        except (EbmlException), e:
            print >>sys.stderr, 'Failed to write integer', e
        
    def _writeSize(self, size):
        bytes = 1
        mask = 0x80
        
        if size != EBML_SIZE_UNKNOWN:
            while ((size >> ((bytes - 1) * 8)) >= (mask - 1) and bytes <= 8):
                mask >>= 1
                bytes += 1
                
            # check for invalid size
            if bytes > 8:
                # warning: invalid size
                mask = 0x01
                bytes = 8
                size = EBML_SIZE_UNKNOWN
        else:
            bytes = 8
            mask = 0x01
            
        try:
            self.write(self._toWriteSize(size, bytes, mask))
        except (EbmlException), e:
            print >>sys.stderr, 'Failed to Write Size', e
            
        
    def _writeInt(self, val, size):
        try:
            self.write(self._toWriteInt(val, size))
        except (EbmlException), e:
            print >>sys.stderr, 'Failed to Write Integer', e
        
    def _writeFloat(self, val, size):
        try:
            self.write(self._toWriteFloat(val, size))
        except (EbmlException), e:
            print >>sys.stderr, 'Failed to write Float', e
        
    def _writeString(self, val):
        self.write(val)
        
    def _writeBinary(self, val):
        self.write(val)
            
    def _getUintSize(self, num):
        size = 1
        
        while (num >= (1 << (size * 8)) and size < 8):
            size += 1
            
        return size  
    
    def _getLengthAndMask(self, b):
        if b & 0x80:
            return (1, 0x80)
        elif b & 0x40:
            return (2, 0x40)
        elif b & 0x20:
            return (3, 0x20)
        elif b & 0x10:
            return (4, 0x10)
        elif b & 0x08:
            return (5, 0x08)
        elif b & 0x04:
            return (6, 0x04)
        elif b & 0x02:
            return (7, 0x02)
        elif b & 0x01:
            return (8, 0x01)
        else:
            # dodgy number
            print >>sys.stderr, 'Unable to get Length of Size with leading 0x%X' % b,
            return (8, 0x01)
        
    def startContainer(self, id):
        # start a container of ebml
        c = EbmlContainer(id)
        self.containers.append(c)
        
        c.position = self.tell()
        
        if not self.useBuffer:        
            self._writeID(id)
            self._writeSize(EBML_SIZE_UNKNOWN)
            c.dataPosition = self.tell()
        
    def endContainer(self):
        # close a previously opened container
        try:
            c = self.containers.pop()
        except:
            print >>sys.stderr, 'Failed to close container, something went very wrong'
            return
        
        if self.useBuffer:
            # write the ID, size and buffer out
            self._writeID(c.id)
            self._writeSize(c.size())
            self.write(c.data)
        else:
            # update the previously written size
            cpos = self.tell()
            self.updateSize(c.position, cpos - c.dataPosition)
            
        del c
        
    # pos = the location of the ID to replace the size of
    def updateSize(self, pos, size):
        cpos = self.tell()
        self.seek(pos)
        
        self.readID()
        
        npos = self.tell()
        
        # get the size and mask of the current size, update
        b1 = self.read(1)
        b1b = ord(b1)
        alen, mask = self._getLengthAndMask(b1b)
        
        self.seek(npos)
        
        # write the new size
        self.write(self._toWriteSize(size, alen, mask))
        
        # shift back to current location
        self.seek(cpos)
        
    # here pos = the location of the ID to replace
    # note currently it uses the existing size, if it is bigger it will fail
    def updateIntValue(self, pos, val):
        cpos = self.tell()
        self.seek(pos)
        
        self.readID()
        size = self.readSize()
        
        # write the new data at pos
        self._writeInt(val, size)
        
        # return to current location
        self.seek(cpos) 
        
    def writeInt(self, id, val):
        
        asize = self._getUintSize(val)
                
        self._writeID(id)
        self._writeSize(asize)
        self._writeInt(val, asize)
        
    def writeFloat(self, id, val):
        self._writeID(id)
        self._writeSize(8)
        self._writeFloat(val, 8)
        
    def writeString(self, id, val):
        nv = str(val)
        alen = len(nv)
        
        self._writeID(id)
        self._writeSize(alen)
        self._writeString(nv)
    
    def writeUTF8(self, id, val):
        try:
            nv = val.encode( 'utf-8' )
        except:
            nv = val

        alen = len(nv)        
        
        self._writeID(id)
        self._writeSize(alen)
        self._writeBinary(nv)
        
    def writeBinary(self, id, val):
        alen = len(val)
        
        self._writeID(id)
        self._writeSize(alen)
        self._writeBinary(val)
        
        
     
        